﻿using System.Collections.Generic;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Profiling;

namespace Mars.GpuInstance
{
    public class GPUAnimatorInstancerRuntimeData
    {
        public Dictionary<int, GPUAnimationClipData> animationClipDataDict;
        public Dictionary<int, int> clipHashToIndex;

        /// <summary>
        /// <para>index: 0 x -> frameNo1, y -> frameNo2, z -> frameNo3, w -> frameNo4</para> 
        /// <para>index: 1 x -> weight1, y -> weight2, z -> weight3, w -> weight4</para> 
        /// </summary>
        public NativeList<float4> animationData;
        
        /// <summary>
        /// 0 to 4: x ->  minFrame, y -> maxFrame (小于0表示非循环动画), z -> speed, w -> startTime
        /// </summary>
        public NativeList<float4> animatorControllerData;

        public List<GPUAnimator> transitioningAnimators;
        
        public bool animationDataModified;
        public bool animatorDataModified;

        public NativeList<float4x4> input;
        public NativeList<float4> campColorData;    //xyz表示颜色，w表示透明度
        public NativeList<bool> inputNonPropScaling;     //是否需要非等比缩放动画
        
        public NativeParallelHashMap<int, int> idToIndexMap; //id和序号映射
        public NativeParallelHashMap<int, int> indexToIdMap; //序号和id映射
        
        protected MaterialPropertyBlock mpb;

        public int totalBoneCount { private set; get; }
        
        //数据是否有更新
        public bool dataModified = true;
        

        private NativeArray<AnchorInfo> anchorInfos;
        private Dictionary<int, int> anchorNameToBoneIndex;
        
        
        public List<GPURenderBatch> skinRenderBatches = new List<GPURenderBatch>();
        public List<GPURenderBatch> shadowBatches = new List<GPURenderBatch>();
        
        public int clearTimeFrame = 1800;  //如果数量清空了，60s后清除

        public ShadowLod animatorShadowLod = ShadowLod.AUTO_SELECT;
        
        private static Vector4[] clipBoneDataArr = new Vector4[1024];              // 最大骨骼数量64，最大动画数量16

        public bool drawable = true;

        public float2 impostorShowRange = -1;

        public int frameRate;
        public int textureSizeX;
        
        public GPUAnimatorInstancerRuntimeData()
        {
            transitioningAnimators = new List<GPUAnimator>();
            anchorNameToBoneIndex = new Dictionary<int, int>(30);
            InitContainer();
            
            clearTimeFrame = 30 * 60 * 5;
        }

        public void InitContainer()
        {
            if(!input.IsCreated) input = new NativeList<float4x4>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
            if(!campColorData.IsCreated) campColorData = new NativeList<float4>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
            if(!inputNonPropScaling.IsCreated) inputNonPropScaling = new NativeList<bool>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
            if(!idToIndexMap.IsCreated) idToIndexMap = new NativeParallelHashMap<int, int>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
            if(!indexToIdMap.IsCreated) indexToIdMap = new NativeParallelHashMap<int, int>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
            if(!animationData.IsCreated) animationData = new NativeList<float4>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
            if(!animatorControllerData.IsCreated) animatorControllerData = new NativeList<float4>(InstancerConstants.ADD_COUNT, Allocator.Persistent);
        }

        public void AddGPUSkinRenderBatch(GPURenderBatch batch)
        {
            skinRenderBatches.Add(batch);
        }

        public void AddShadowBatch(GPURenderBatch batch)
        {
            shadowBatches.Add(batch);
        }
        
        public void UpdateMpb(MaterialPropertyBlock _mpb)
        {
            mpb = _mpb;
        }
        
        public void UpdateTotalBoneCount(int totalBoneCount)
        {
            this.totalBoneCount = totalBoneCount;
            mpb.SetFloat(InstancerConstants.TOTAL_NUMBER_OF_BONES, totalBoneCount);
        }

        public void UpdateFrameRate(float frameRate)
        {
            mpb.SetFloat(InstancerConstants.FRAME_RATE, frameRate);
        }

        public void UpdateAnimationTex(Texture2D animTex)
        {
            if (animTex)
            {
                mpb.SetTexture(InstancerConstants.ANIM_TEX, animTex);
            }
        }
        
        public void UpdateClipBoneData(GPUAnimationData animationData, string name)
        {
            var clipDataList = animationData.clipDataList;
            var clipCount = clipDataList.Count;
            clipBoneDataArr[0] = new Vector4(clipCount, 0, 0, 0);   // 第一个元素存储clip的数量
            
            var dataCount = 1 + clipCount + clipCount * animationData.totalBoneCount / 2;
            if (dataCount > clipBoneDataArr.Length)
            {
                Debug.LogError(
                    $"[{nameof(GPURenderBatch)}] 动画骨骼超出最大限制，请联系 Kaze 支持。 name = {name},  clipCount = {clipCount}, boneCount = {animationData.totalBoneCount}, totalDataCount = {dataCount}");
            }
            
            for (int i = 0; i < clipCount; i++)
            {
                var clipData = clipDataList[i];

                var arrIndex = i + 1;
                var gpuClipData = clipBoneDataArr[arrIndex];
                gpuClipData.x = clipData.animTexDataStartIndex;
                gpuClipData.y = clipData.animTexNonPropDataStartIndex;
                gpuClipData.z = clipData.propBoneCount;
                clipBoneDataArr[arrIndex] = gpuClipData;

                for (int j = 0; j < animationData.totalBoneCount; j++)
                {
                    var boneData = clipData.boneDataArr[j];
                    
                    arrIndex = (1 + clipCount) + (i * animationData.totalBoneCount + j) / 2;
                    if (arrIndex >= clipBoneDataArr.Length) continue;
                    var channelIndex = (i * animationData.totalBoneCount + j) * 2 % 4;
                    gpuClipData = clipBoneDataArr[arrIndex];
                    gpuClipData[channelIndex] = boneData.nonPropScaling ? 1 : 0;
                    gpuClipData[channelIndex + 1] = boneData.index;
                    clipBoneDataArr[arrIndex] = gpuClipData;
                }
            }
            
            mpb.SetVectorArray(InstancerConstants.GPU_ClipBoneData, clipBoneDataArr);
        }

        #region 挂点相关

        //更新挂点信息
        public void UpdateAnchorInfos(NativeArray<AnchorInfo> _anchorInfos)
        {
            anchorInfos = _anchorInfos;
        }

        public void AddAnchorNameToBoneIndex(int nameHash, int boneIndex)
        {
            anchorNameToBoneIndex.TryAdd(nameHash, boneIndex);
        }
        
        public void UpdateAnchorNameToBoneIndex(Dictionary<int, int> _anchorNameToBoneIndex)
        {
            anchorNameToBoneIndex = _anchorNameToBoneIndex;
        }
        
        //获取挂点信息
        public AnchorInfo GetAnchorInfoByBoneIndex(int boneIndex, int index, int frameRate = 30)
        {
            if (anchorInfos.IsCreated)
            {
                int maxWeightIndex = 0;
                var animationWeightData = animationData[index * 2 + 1];
                float maxWeight = animationWeightData[0];
                for (int i = 1; i < 2; i++)
                {
                    if (animationWeightData[i] > maxWeight)
                    {
                        maxWeight = animationWeightData[i];
                        maxWeightIndex = i;
                    }
                }
                
                var crowdAnimatorData = animatorControllerData[index * 2 + maxWeightIndex];
                if (!float.IsNaN(crowdAnimatorData.w))
                {
                    var currentFrame = crowdAnimatorData.x;
                    float MIN_SPEED = 0.000001f;
                
                    int isLooping = (int) math.sign(crowdAnimatorData.y); // y == 0, 返回0， y>0返回1，y < 0 返回 -1
                    crowdAnimatorData.y = crowdAnimatorData.y * isLooping; //设置成绝对值
                    float clipTime = math.abs(Time.time - crowdAnimatorData.w) * crowdAnimatorData.z;
                    float clipFrameCount = crowdAnimatorData.y - crowdAnimatorData.x;
                    float length = clipFrameCount / frameRate;

                    if (crowdAnimatorData.z > MIN_SPEED)
                    {
                        if (isLooping < 0 && clipTime > length) //非循环动画
                        {
                            currentFrame = crowdAnimatorData.y;
                        }
                        else
                        {
                            //frac取小数部分
                            currentFrame = math.frac(clipTime / length) * clipFrameCount + crowdAnimatorData.x;
                        }
                    }
                
                    int frameIndex = (int)math.floor(currentFrame) * totalBoneCount + boneIndex;
                   // Debug.LogWarning($"[see see] curtime: {crowdAnimatorData.w} frameIndex: {frameIndex} boneIndex: {boneIndex} currentFrame: {currentFrame} anchorInfos.Length: {anchorInfos.Length}");
                
                    if (anchorInfos.Length > frameIndex && frameIndex >= 0)
                    {
                        return anchorInfos[frameIndex];
                    }
                }
            }
            return default;
        }
        
        //获取挂点信息
        public AnchorInfo GetAnchorInfoByAnchorNameHash(int nameHash, int index, int frameRate = 30)
        {
            if (anchorNameToBoneIndex.TryGetValue(nameHash, out int boneIndex))
            {
                return GetAnchorInfoByBoneIndex(boneIndex, index, frameRate);
            }

            return default;
        }

        #endregion
      
        

        /// <summary>
        /// 获取实体的序号，小于0则表示不存在
        /// </summary>
        /// <param name="id"></param>
        /// <returns></returns>
        public int GetInstanceIndex(int id)
        {
            if (idToIndexMap.TryGetValue(id, out var index))
            {
                return index;
            }

            return -1;
        }

        /// <summary>
        /// 获取实体的id，小于0则表示不存在
        /// </summary>
        /// <param name="index"></param>
        /// <returns></returns>
        public int GetInstanceId(int index)
        {
            if (indexToIdMap.TryGetValue(index, out var id))
            {
                return id;
            }
            return -1;
        }



        public bool AddInstance(int id, float4x4 matrix)
        {
            dataModified = true;
            animationDataModified = true;
            animatorDataModified = true;
            clearTimeFrame = int.MaxValue;

            if (idToIndexMap.TryGetValue(id, out var index))
            {
                input[index] = matrix;
                return false;
            }
            else
            {
                idToIndexMap.TryAdd(id, input.Length);
                indexToIdMap.TryAdd(input.Length, id);
                input.Add(matrix);
                inputNonPropScaling.Add(false);
                campColorData.Add(new float4(1, 1, 1, 1));
                InitAnimationData();
                return true;
            }
        }

        public bool AddInstance(int id,  Vector3 pos, Vector3 scale, Quaternion rot)
        {
            return AddInstance(id, Matrix4x4.TRS(pos, rot, scale));
        }
        

        /// <summary>
        /// 移除实体，把最后一个跟当前的交换
        /// </summary>
        /// <param name="id"></param>
        public void RemoveInstance(int id)
        {
            if (!idToIndexMap.TryGetValue(id, out var index))
            {
                return;
            }
            
            Profiler.BeginSample("GPU RemoveInstance");

            int lastIndex = input.Length - 1;
            if (indexToIdMap.TryGetValue(lastIndex, out var lastId)) //找到最后一个
            {
                input[index] = input[^1];   //把最后一个替换到当前index
                inputNonPropScaling[index] = inputNonPropScaling[^1];
                campColorData[index] = campColorData[^1];
                input.RemoveAt(lastIndex);  //把最后一个删掉
                inputNonPropScaling.RemoveAt(lastIndex);
                campColorData.RemoveAt(lastIndex);
                RemoveAnimationData(index); //把动画数据删除，同样把最后一个和当前的交换，然后删除最后一个
                

                idToIndexMap[lastId] = index;
                indexToIdMap[index] = lastId;
                
                idToIndexMap.Remove(id);
                indexToIdMap.Remove(lastIndex);
                
                dataModified = true;
                animationDataModified = true;
                animatorDataModified = true;
                
                var len = transitioningAnimators.Count;
                for (int i = len - 1; i >= 0; i--)
                {
                    var animator = transitioningAnimators[i];
                    if (animator.transition != null)
                    {
                        if (animator.transition.arrayIndex == index)
                        {
                            transitioningAnimators.RemoveAt(i);
                            break;
                        }
                    }
                }

                len = transitioningAnimators.Count;
                for (int i = 0; i < len; i++)
                {
                    var animator = transitioningAnimators[i];
                    animator.transitionIndex = i;
                    if (animator.transition != null && animator.transition.arrayIndex == lastIndex)
                    {
                        animator.transition.UpdateArrayIndex(index);
                    }
                }

                //记录清空时的帧数
                if (input.Length <= 0)
                {
                    clearTimeFrame = Time.frameCount;
                }
            }
            
            Profiler.EndSample();
        }
        
        public void ClearInstance()
        {
            input.Clear();
            campColorData.Clear();
            inputNonPropScaling.Clear();
            idToIndexMap.Clear();
            indexToIdMap.Clear();
            dataModified = true;
            animationData.Clear();
            animatorControllerData.Clear();
            animationDataModified = true;
            animatorDataModified = true;
            
            //记录清空时的帧数
            if (input.Length <= 0)
            {
                clearTimeFrame = Time.frameCount;
            }
        }

        public void Dispose()
        {
            if (input.IsCreated)
            {
                input.Dispose();
            }

            if (campColorData.IsCreated)
            {
                campColorData.Dispose();
            }

            if (inputNonPropScaling.IsCreated)
            {
                inputNonPropScaling.Dispose();
            }
            
            if (idToIndexMap.IsCreated)
            {
                idToIndexMap.Dispose();
            }

            if (indexToIdMap.IsCreated)
            {
                indexToIdMap.Dispose();
            }

            if (animatorControllerData.IsCreated)
            {
                animatorControllerData.Dispose();
            }

            if (animationData.IsCreated)
            {
                animationData.Dispose();
            }
            

            foreach (var batch in skinRenderBatches)
            {
                batch.Dispose();
            }
            skinRenderBatches.Clear();

            foreach (var batch in shadowBatches)
            {
                batch.Dispose();
            }
            shadowBatches.Clear();
            
            transitioningAnimators.Clear();
            
            animationClipDataDict?.Clear();
            clipHashToIndex?.Clear();
            
            if (mpb != null)
            {
                mpb.Clear();
                mpb = null;
            }
        }

        public void VisibilityCalculation(NativeArray<float4> frustumPlanes, float3 cameraPos, NativeList<JobHandle> handles)
        {
            if (input.Length <= 0)
            {
                animationDataModified = true;
                animatorDataModified = true;
                dataModified = false;
                return;
            }
            
            foreach (var batch in skinRenderBatches)
            {
                batch.VisibilityCalculation(this, frustumPlanes, cameraPos, handles);
            }
            
            foreach (var batch in shadowBatches)
            {
                batch.VisibilityCalculation(this, frustumPlanes, cameraPos, handles);
            }

            animationDataModified = true;
            animatorDataModified = true;
            dataModified = false;
        }

        public void CopyAnimationData(NativeList<JobHandle> handles)
        {
            if (input.Length <= 0)
            {
                animationDataModified = false;
                animatorDataModified = false;
                return;
            }
            
            foreach (var batch in skinRenderBatches)
            {
                batch.CopyAnimationData( animatorControllerData, animationData, handles);
            }
            
            foreach (var batch in shadowBatches)
            {
                batch.CopyAnimationData( animatorControllerData, animationData, handles);
            }
            
            animationDataModified = false;
            animatorDataModified = false;
        }


        public void SubmitToGPU()
        {
            foreach (var batch in skinRenderBatches)
            {
                batch.SubmitToGPU();
            }
            
            foreach (var batch in shadowBatches)
            {
                batch.SubmitToGPU();
            }
        }


        /// <summary>
        /// 初始化动画数据，
        /// 每个实体对于两个animationData，2个animatorControllerData
        /// </summary>
        private void InitAnimationData()
        {
            for (int i = 0; i < 2; i++)
            {
                animationData.Add(float4.zero);
                animatorControllerData.Add(float4.zero);
            }
        }

        /// <summary>
        /// 删除动画数据
        /// animationData删除两个，animatorControllerData删除4个
        /// </summary>
        /// <param name="index"></param>
        private void RemoveAnimationData(int index)
        {
            int animatorIndex = index * 2;
            animationData[animatorIndex] = animationData[^2];
            animationData[animatorIndex + 1] = animationData[^1];
            animationData.RemoveRange(animationData.Length - 2, 2);

            animatorControllerData[animatorIndex] = animatorControllerData[^2];
            animatorControllerData[animatorIndex + 1] = animatorControllerData[^1];
            animatorControllerData.RemoveRange(animatorControllerData.Length - 2, 2);

            animationDataModified = true;
            animatorDataModified = true;
        }
        
        public void CalculateAnim(NativeList<JobHandle> handles)
        {
            if (input.Length <= 0)
            {
                return;
            }
            
            foreach (var batch in skinRenderBatches)
            {
                var handle = batch.CalculateAnim(frameRate, textureSizeX, totalBoneCount);
                handles.Add(handle);
            }
            
            foreach (var batch in shadowBatches)
            {
                var handle = batch.CalculateAnim(frameRate, textureSizeX, totalBoneCount);
                handles.Add(handle);
            }
        }
        
        #region Draw

        public void Draw(Bounds drawBound, bool drawShadow = true)
        {
            if (!drawable)
            {
                return;
            }
            if (input.Length <= 0)
            {
                return;
            }
            
            foreach (var batch in skinRenderBatches)
            {
                batch.Draw(drawBound, animatorShadowLod);
            }

            if (animatorShadowLod != ShadowLod.FOLLOW_LOD)
            {
                if ( (drawShadow && animatorShadowLod != ShadowLod.DISABLE))
                {
                    foreach (var batch in shadowBatches)
                    {
                        batch.DrawShadows(drawBound);
                    }
                }
            }
        }

        public void DrawDebug(float sightDistance)
        {
            if (input.Length <= 0)
            {
                return;
            }
            
            foreach (var batch in skinRenderBatches)
            {
                batch.DrawDebug(sightDistance);
            }
        }

        #endregion

        public void OnRecycle()
        {
            if (input.IsCreated)
            {
                input.Dispose();
            }

            if (campColorData.IsCreated)
            {
                campColorData.Dispose();
            }

            if (inputNonPropScaling.IsCreated)
            {
                inputNonPropScaling.Dispose();
            }
            
            if (idToIndexMap.IsCreated)
            {
                idToIndexMap.Dispose();
            }

            if (indexToIdMap.IsCreated)
            {
                indexToIdMap.Dispose();
            }
            
            if (animatorControllerData.IsCreated)
            {
                animatorControllerData.Dispose();
            }

            if (animationData.IsCreated)
            {
                animationData.Dispose();
            }

            var cacheMgr = GPUAnimatorInstanceCacheDataManager.Instance;

            foreach (var batch in skinRenderBatches)
            {
                if (batch is GPUSkinRenderBatch renderBatch)
                {
                    cacheMgr.RecycleGPUSkinRenderBatch(renderBatch);
                }
                else if (batch is GPUMeshRenderBatch meshRenderBatch)
                {
                    cacheMgr.RecycleGPUMeshRenderBatch(meshRenderBatch);
                }
            }
            skinRenderBatches.Clear();

            foreach (var batch in shadowBatches)
            {
                if (batch is GPUSkinRenderBatch renderBatch)
                {
                    cacheMgr.RecycleGPUSkinRenderBatch(renderBatch);
                }
                else if (batch is GPUMeshRenderBatch meshRenderBatch)
                {
                    cacheMgr.RecycleGPUMeshRenderBatch(meshRenderBatch);
                }
            }
            shadowBatches.Clear();

            foreach (var animator in transitioningAnimators)
            {
                cacheMgr.RecycleGPUAnimator(animator);
            }
            transitioningAnimators.Clear();
            
            animationClipDataDict.Clear();
            clipHashToIndex.Clear();
            
            anchorNameToBoneIndex.Clear();

            if (mpb != null)
            {
                mpb.Clear();
                mpb = null;
            }
        }
    }
}
